#include "cuSlideWinCropImg.h"

void calcNumCrops(int src_h, int src_w, int win_h, int win_w, double ratio_h, double ratio_w, int& num_h, int& num_w)
{
    int step_h = int(win_h * ratio_h);
    int step_w = int(win_w * ratio_w);
    num_h = int(ceil((double)(src_h - win_h + 1) / (double)step_h));
    num_w = int(ceil((double)(src_w - win_w + 1) / (double)step_w));
}

__global__ void cuSlideWinCropImg_kernel(uint8_t* src, uint8_t* dst, int src_h, int src_w, int win_h, int win_w, int num_h, int num_w, int step_h, int step_w)
{
    int tx = threadIdx.x, ty = threadIdx.y;
    int bx = blockIdx.x, by = blockIdx.y, bz = blockIdx.z;

    // 计算目前block处于哪个crop
    int idx_w = bz % num_w, idx_h = bz / num_w;

    // 找到src/dst目前crop的首指针位置
    uchar3* dst_u3_cur = (uchar3*)(dst) + bz * win_h * win_w;
    uchar3* src_u3_cur = (uchar3*)(src) + (idx_h * step_h * src_w) + (idx_w * step_w);

    int x = tx + bx * blockDim.x;
    int y = ty + by * blockDim.y;

    // if (x < win_w && y < win_h && idx_w * step_w + x < src_w && idx_h * step_h + y < src_h)
    if (x < win_w && y < win_h)//最后不足一个window尺寸的crop会被舍弃，因此不会越界
    {
        dst_u3_cur[y * win_w + x] = src_u3_cur[y * src_w + x];

        // if (bz == 0 && x == 0 && y == 0)
        // {
        //     printf("src: %d\n", int(src_u3_cur[y * src_w + x].x));
        //     printf("src: %d\n", int(src_u3_cur[y * src_w + x].y));
        //     printf("src: %d\n", int(src_u3_cur[y * src_w + x].z));
        //     printf("dst: %d\n", int(dst_u3_cur[y * win_w + x].x));
        //     printf("dst: %d\n", int(dst_u3_cur[y * win_w + x].y));
        //     printf("dst: %d\n", int(dst_u3_cur[y * win_w + x].z));
        // }
    }
}

cudaError_t cuSlideWinCropImg_launch(uint8_t* src, uint8_t* dst, int src_h, int src_w, int win_h, int win_w, double ratio_h, double ratio_w, int num_h, int num_w, cudaStream_t stream)
{
    if (!src || !dst)
    return cudaErrorInvalidDevicePointer;

    // block 的大小设置成 warp 的大小
    const dim3 blockDim(BLOCK_SIZE_16, BLOCK_SIZE_16, 1);//blockDim最好为32的整数倍，否则会造成线程浪费
    const dim3 gridDim(iDivUp(win_w, blockDim.x), iDivUp(win_h, blockDim.y), num_w * num_h);

    int step_h = int(win_h * ratio_h);
    int step_w = int(win_w * ratio_w);

    // using namespace std;
    // cout << "src_h: " << src_h << endl;
    // cout << "src_w: " << src_w << endl;
    // cout << "win_h: " << win_h << endl;
    // cout << "win_w: " << win_w << endl;
    // cout << "num_h: " << num_h << endl;
    // cout << "num_w: " << num_w << endl;
    // cout << "step_h: " << step_h << endl;
    // cout << "step_w: " << step_w << endl;

    cuSlideWinCropImg_kernel << <gridDim, blockDim, 0, stream >> > (src, dst, src_h, src_w, win_h, win_w, num_h, num_w, step_h, step_w);

    return CUDA(cudaGetLastError());
}
